{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import necessary modules\n",
    "# uncomment to get plots displayed in notebook\n",
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from classy import Class\n",
    "from scipy.optimize import fsolve\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "############################################\n",
    "#\n",
    "# Varying parameter (others fixed to default)\n",
    "#\n",
    "var_name = 'N_ur'\n",
    "var_array = np.linspace(3.044,5.044,5)\n",
    "var_num = len(var_array)\n",
    "var_legend = r'$N_\\mathrm{eff}$'\n",
    "var_figname = 'neff'\n",
    "#\n",
    "# Constraints to be matched\n",
    "#\n",
    "# As explained in the \"Neutrino cosmology\" book, CUP, Lesgourgues et al., section 5.3, the goal is to vary\n",
    "# - omega_cdm by a factor alpha = (1 + coeff*Neff)/(1 + coeff*3.046)\n",
    "# - h by a factor sqrt*(alpha)\n",
    "# in order to keep a fixed z_equality(R/M) and z_equality(M/Lambda)\n",
    "#\n",
    "omega_b = 0.0223828\n",
    "omega_cdm_standard = 0.1201075\n",
    "h_standard = 0.67810\n",
    "#\n",
    "# coefficient such that omega_r = omega_gamma (1 + coeff*Neff),\n",
    "# i.e. such that omega_ur = omega_gamma * coeff * Neff:\n",
    "# coeff = omega_ur/omega_gamma/Neff_standard \n",
    "# We could extract omega_ur and omega_gamma on-the-fly within th script, \n",
    "# but for simplicity we did a preliminary interactive run with background_verbose=2\n",
    "# and we copied the values given in the budget output.\n",
    "#\n",
    "coeff = 1.70961e-05/2.47298e-05/3.044\n",
    "print (\"coeff=\",coeff)\n",
    "#\n",
    "#############################################\n",
    "#\n",
    "# Fixed settings\n",
    "#\n",
    "common_settings = {# fixed LambdaCDM parameters\n",
    "                   'omega_b':omega_b,\n",
    "                   'A_s':2.100549e-09,\n",
    "                   'n_s':0.9660499,\n",
    "                   'tau_reio':0.05430842,\n",
    "                   # output and precision parameters\n",
    "                   'output':'tCl,pCl,lCl,mPk',\n",
    "                   'lensing':'yes',\n",
    "                   'P_k_max_1/Mpc':3.0,\n",
    "                   'l_switch_limber':9}  \n",
    "#\n",
    "##############################################\n",
    "#\n",
    "# loop over varying parameter values\n",
    "#\n",
    "M = {}\n",
    "#\n",
    "for i, N_ur in enumerate(var_array):\n",
    "    #\n",
    "    # rescale omega_cdm and h\n",
    "    #\n",
    "    alpha = (1.+coeff*N_ur)/(1.+coeff*3.044)\n",
    "    omega_cdm = (omega_b + omega_cdm_standard)*alpha - omega_b\n",
    "    h = h_standard*math.sqrt(alpha)\n",
    "    print (' * Compute with %s=%e, %s=%e, %s=%e'%('N_ur',N_ur,'omega_cdm',omega_cdm,'h',h))\n",
    "    #\n",
    "    # call CLASS\n",
    "    #\n",
    "    M[i] = Class()\n",
    "    M[i].set(common_settings)\n",
    "    M[i].set({'N_ur':N_ur})\n",
    "    M[i].set({'omega_cdm':omega_cdm})\n",
    "    M[i].set({'h':h})\n",
    "    M[i].compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# esthetic definitions for the plots\n",
    "font = {'size'   : 24, 'family':'STIXGeneral'}\n",
    "axislabelfontsize='large'\n",
    "matplotlib.rc('font', **font)\n",
    "matplotlib.mathtext.rcParams['legend.fontsize']='medium'\n",
    "plt.rcParams[\"figure.figsize\"] = [8.0,6.0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#############################################\n",
    "#\n",
    "# extract spectra and plot them\n",
    "#\n",
    "#############################################\n",
    "kvec = np.logspace(-4,np.log10(3),1000) # array of kvec in h/Mpc\n",
    "twopi = 2.*math.pi\n",
    "#\n",
    "# Create figures\n",
    "#\n",
    "fig_Pk, ax_Pk = plt.subplots()\n",
    "fig_TT, ax_TT = plt.subplots()\n",
    "#\n",
    "# loop over varying parameter values\n",
    "#\n",
    "ll = {}\n",
    "clM = {}\n",
    "clTT = {}\n",
    "pkM = {}\n",
    "legarray = []\n",
    "\n",
    "for i, N_ur in enumerate(var_array):\n",
    "    #\n",
    "    alpha = (1.+coeff*N_ur)/(1.+coeff*3.044)\n",
    "    h = 0.67810*math.sqrt(alpha) # this is h\n",
    "    #\n",
    "    # deal with colors and legends\n",
    "    #\n",
    "    if i == 0:\n",
    "        var_color = 'k'\n",
    "        var_alpha = 1.\n",
    "    else:\n",
    "        var_color = plt.cm.Reds(0.8*i/(var_num-1))\n",
    "    #\n",
    "    # get Cls\n",
    "    #\n",
    "    clM[i] = M[i].lensed_cl(2500)\n",
    "    ll[i] = clM[i]['ell'][2:]\n",
    "    clTT[i] = clM[i]['tt'][2:]\n",
    "    #\n",
    "    # store P(k) for common k values\n",
    "    #\n",
    "    pkM[i] = []\n",
    "    # The function .pk(k,z) wants k in 1/Mpc so we must convert kvec for each case with the right h \n",
    "    khvec = kvec*h # This is k in 1/Mpc\n",
    "    for kh in khvec:\n",
    "        pkM[i].append(M[i].pk(kh,0.)*h**3) \n",
    "    #    \n",
    "    # plot P(k)\n",
    "    #\n",
    "    if i == 0:\n",
    "        ax_Pk.semilogx(kvec,np.array(pkM[i])/np.array(pkM[0]),\n",
    "                       color=var_color,#alpha=var_alpha,\n",
    "                       linestyle='-')\n",
    "    else:\n",
    "        ax_Pk.semilogx(kvec,np.array(pkM[i])/np.array(pkM[0]),\n",
    "                       color=var_color,#alpha=var_alpha,\n",
    "                       linestyle='-',\n",
    "                      label=r'$\\Delta N_\\mathrm{eff}=%g$'%(N_ur-3.044))\n",
    "    #\n",
    "    # plot C_l^TT\n",
    "    #\n",
    "    if i == 0:\n",
    "        ax_TT.semilogx(ll[i],clTT[i]/clTT[0],\n",
    "                       color=var_color,alpha=var_alpha,linestyle='-')\n",
    "    else:    \n",
    "        ax_TT.semilogx(ll[i],clTT[i]/clTT[0],\n",
    "                       color=var_color,alpha=var_alpha,linestyle='-',\n",
    "                      label=r'$\\Delta N_\\mathrm{eff}=%g$'%(N_ur-3.044))\n",
    "#\n",
    "# output of P(k) figure\n",
    "#\n",
    "ax_Pk.set_xlim([1.e-3,3.])\n",
    "ax_Pk.set_ylim([0.98,1.20])\n",
    "ax_Pk.set_xlabel(r'$k \\,\\,\\,\\, [h^{-1}\\mathrm{Mpc}]$')\n",
    "ax_Pk.set_ylabel(r'$P(k)/P(k)[N_\\mathrm{eff}=3.046]$')\n",
    "ax_Pk.legend(loc='upper left')\n",
    "fig_Pk.tight_layout()\n",
    "fig_Pk.savefig('ratio-%s-Pk.pdf' % var_figname)\n",
    "#\n",
    "# output of C_l^TT figure\n",
    "#      \n",
    "ax_TT.set_xlim([2,2500])\n",
    "ax_TT.set_ylim([0.850,1.005])\n",
    "ax_TT.set_xlabel(r'$\\mathrm{Multipole} \\,\\,\\,\\,  \\ell$')\n",
    "ax_TT.set_ylabel(r'$C_\\ell^\\mathrm{TT}/C_\\ell^\\mathrm{TT}(N_\\mathrm{eff}=3.046)$')\n",
    "ax_TT.legend(loc='lower left')\n",
    "fig_TT.tight_layout()\n",
    "fig_TT.savefig('ratio-%s-cltt.pdf' % var_figname)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
